import pandas as pd
import matplotlib.pyplot as plt
import os

df = pd.read_excel('VIXCLS.xls', skiprows=10, index_col='observation_date')
df = df[df.VIXCLS != 0] # the "#N/A"s are imported as 0

def my_plotter(df, column, filename, save=True, annotate=False):
    """
    Plots VIX data, with options for annotations.
    Original data from St. Louis Fed website
    Returns plots as PDF, SVG, and EPS files.
    """
    fig, ax = plt.subplots(1, 1)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.set_ylabel('VIX', fontsize=24)
    plt.plot(df.index, df[column], label='VIX', color='blue')
    if annotate:
        crisis = df.truncate(after=pd.Timestamp('2010-01-01'))
        crisis_peak = crisis.sort_values(by='VIXCLS', ascending=False).head(1)
        plt.annotate('Financial Crisis', xy=(crisis_peak.index, crisis_peak['VIXCLS'][0]))
        covid = df.truncate(before=pd.Timestamp('2010-01-01'))
        covid_peak = covid.sort_values(by='VIXCLS', ascending=False).head(1)
        plt.annotate('Covid-19', xy=(covid_peak.index, covid_peak['VIXCLS'][0]))
    if save and not annotate:
        plt.savefig(f'{filename}_raw.pdf')
        plt.savefig(f'{filename}_raw.svg')
        plt.savefig(f'{filename}_raw.eps')
    if save and annotate:
        plt.savefig(f'{filename}.pdf')
        plt.savefig(f'{filename}.svg')
        plt.savefig(f'{filename}.eps')



if __name__ == '__main__':
    my_plotter(df, 'VIXCLS', 'fig1', save=True, annotate=True)
    # Use next line to produce figure 1 without annotaions
    # my_plotter(df, 'VIXCLS', 'fig1', save=True)
